"""
蛋白质搜索工具
调用后端API /api/search/ 进行蛋白质搜索
"""

import json
import requests
from qwen_agent.tools import BaseTool
from typing import Dict, Any
import config


class ProteinSearchTool(BaseTool):
    """蛋白质搜索工具类"""
    
    # 必须设置的类属性
    name = "protein_search"
    description = "根据关键词、物种、类型等条件搜索蛋白质数据库。可以搜索蛋白质名称、登录号(accession)、物种名称、蛋白质类型等信息。"
    parameters = {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "搜索关键词，将在物种名称、登录号(accession)、蛋白质类型中搜索"
            },
            "species": {
                "type": "string",
                "description": "按物种名称过滤，例如'E. coli'、'Escherichia coli'、'Homo sapiens'等"
            },
            "protein_type": {
                "type": "string",
                "description": "按蛋白质类型过滤，例如'AcrIIA'、'CRISPR'、'Cas'等"
            },
            "page": {
                "type": "integer",
                "description": "页码，从1开始（默认1）",
                "default": 1
            },
            "page_size": {
                "type": "integer",
                "description": "每页返回的结果数量（默认20，最大100）",
                "default": 20
            }
        },
        "required": ["query"]
    }
    
    def call(self, params: Dict[str, Any], **kwargs) -> str:
        """
        执行蛋白质搜索
        
        Args:
            params: 参数字典，包含query、species、protein_type、page、page_size
            **kwargs: 其他关键字参数
            
        Returns:
            str: JSON格式的字符串结果
        """
        try:
            # 参数解析和验证
            if isinstance(params, str):
                params = json.loads(params)
            
            query = params.get("query", "").strip()
            if not query:
                return json.dumps({
                    "success": False,
                    "error": "搜索关键词不能为空"
                }, ensure_ascii=False)
            
            # 获取其他参数
            species = params.get("species", "").strip()
            protein_type = params.get("protein_type", "").strip()
            page = params.get("page", config.DEFAULT_PARAMS["search"]["page"])
            page_size = params.get("page_size", config.DEFAULT_PARAMS["search"]["page_size"])
            
            # 参数验证
            if page < 1:
                page = 1
            if page_size < 1:
                page_size = config.DEFAULT_PARAMS["search"]["page_size"]
            if page_size > config.LIMITS["search"]["max_page_size"]:
                page_size = config.LIMITS["search"]["max_page_size"]
            
            # 构建API请求参数
            api_params = {
                "q": query,
                "page": page,
                "page_size": page_size
            }
            if species:
                api_params["species"] = species
            if protein_type:
                api_params["type"] = protein_type
            
            # 调用API
            api_url = config.API_ENDPOINTS["search"]
            # 添加调试日志
            import logging
            logger = logging.getLogger(__name__)
            logger.info(f"调用蛋白质搜索API: {api_url}, 参数: {api_params}")
            
            response = requests.get(
                api_url,
                params=api_params,
                timeout=30
            )
            response.raise_for_status()
            api_result = response.json()
            
            # 格式化结果
            formatted_result = {
                "success": True,
                "total": api_result.get("total", 0),
                "page": api_result.get("page", page),
                "page_size": api_result.get("page_size", page_size),
                "total_pages": api_result.get("total_pages", 1),
                "query": query,
                "results": []
            }
            
            # 格式化结果列表
            for item in api_result.get("data", []):
                formatted_result["results"].append({
                    "id": item.get("id"),
                    "accession": item.get("accession"),
                    "species": item.get("species"),
                    "type": item.get("type"),
                    "assembly": item.get("assembly"),
                    "has_pdb": item.get("has_pdb", False),
                    "protein_sequence_preview": item.get("protein_sequence_preview", "")
                })
            
            return json.dumps(formatted_result, ensure_ascii=False)
            
        except requests.exceptions.RequestException as e:
            error_msg = {
                "success": False,
                "error": f"API调用失败: {str(e)}"
            }
            return json.dumps(error_msg, ensure_ascii=False)
        
        except Exception as e:
            error_msg = {
                "success": False,
                "error": f"工具执行失败: {str(e)}"
            }
            return json.dumps(error_msg, ensure_ascii=False)

